import os
import argparse
import rosbag
import numpy as np
from bisect import bisect_left
from pyproj import Proj, Transformer

def euler_to_rotmat(roll, pitch, yaw):
    Rx = np.array([
        [1,             0,              0],
        [0, np.cos(roll), -np.sin(roll)],
        [0, np.sin(roll),  np.cos(roll)]
    ])
    Ry = np.array([
        [ np.cos(pitch), 0, np.sin(pitch)],
        [             0, 1,             0],
        [-np.sin(pitch), 0, np.cos(pitch)]
    ])
    Rz = np.array([
        [np.cos(yaw), -np.sin(yaw), 0],
        [np.sin(yaw),  np.cos(yaw), 0],
        [          0,            0, 1]
    ])
    return Rz.dot(Ry.dot(Rx))

def extract_localization_poses(bag_path, euler_in_degrees=False):
    """
    从 rosbag 中读取 /ins/localization，返回两个等长列表：
      - loc_times: [t0, t1, t2, ...]
      - loc_poses: [array(12), array(12), ...]
    其中 t_i = rosbag 读取时所附带的时间戳 t.to_sec()。
    """
    bag = rosbag.Bag(bag_path, "r")
    loc_times = []
    loc_poses = []
    transformer = None
    origin_e = origin_n = origin_alt = None
    cnt = 1

    for topic, msg, t in bag.read_messages(topics=["/ins/localization"]):
        print(cnt)
        cnt = cnt + 1
        # 1) 直接用 bag 提供的 t 作为时间戳
        timestamp = t.to_sec()

        # 2) 提取经纬度/高度
        lat = msg.latitude
        lon = msg.longitude
        alt = msg.altitude

        # 第一次出现就根据该帧经纬度动态生成 UTM 投影 Transformer
        if transformer is None:
            zone_number = int((lon + 180) / 6) + 1
            is_south = (lat < 0)
            proj_utm = Proj(proj='utm', zone=zone_number, datum='WGS84', south=is_south)
            transformer = Transformer.from_proj(
                Proj("epsg:4326"),  # WGS84 经纬度
                proj_utm,
                always_xy=True
            )

        # 3) 计算相对于首帧的平移
        if origin_e is None:
            e0, n0 = transformer.transform(lon, lat)
            origin_e, origin_n, origin_alt = e0, n0, alt

        e, n = transformer.transform(lon, lat)
        x = float(e - origin_e)
        y = float(n - origin_n)
        z = float(alt - origin_alt)

        # 4) 角度转弧度
        if euler_in_degrees:
            roll = np.deg2rad(msg.roll)
            pitch = np.deg2rad(msg.pitch)
            yaw = np.deg2rad(msg.azimuth)
        else:
            roll = msg.roll
            pitch = msg.pitch
            yaw = msg.azimuth

        # 5) 计算 3x3 旋转矩阵
        R = euler_to_rotmat(roll, pitch, yaw)

        # 6) 拼 3x4 并展开成 12 维
        T = np.zeros((3, 4), dtype=np.float64)
        T[:3, :3] = R
        T[0, 3], T[1, 3], T[2, 3] = x, y, z
        T_flat = T.reshape(12)

        loc_times.append(timestamp)
        loc_poses.append(T_flat)

    bag.close()
    return loc_times, loc_poses


def align_with_images(bag_path, loc_times, loc_poses, output_txt):
    """
    把 /camera/image_left 里每张图像的时间戳与 loc_times 对齐，给每帧图像分配最近邻的 pose。
    """
    bag = rosbag.Bag(bag_path, "r")
    image_times = []
    cnt = 0

    for topic, msg, t in bag.read_messages(topics=["/camera/image_left"]):
        print(cnt)
        cnt = cnt + 1
        # 如果 /camera/image_left 的 msg 有 header，就用 msg.header.stamp，否则也可以改用 t.to_sec()
        img_ts = msg.header.stamp.secs + msg.header.stamp.nsecs * 1e-9
        image_times.append(img_ts)

    bag.close()

    # 确保 loc_times 是升序（一般 rosbag 顺序已经排好，这里额外排序以防万一）
    sorted_idx = np.argsort(loc_times)
    loc_times_sorted = [loc_times[i] for i in sorted_idx]
    loc_poses_sorted = [loc_poses[i] for i in sorted_idx]

    from bisect import bisect_left
    with open(output_txt, "w") as fout:
        for img_ts in image_times:
            idx = bisect_left(loc_times_sorted, img_ts)
            if idx == 0:
                best = 0
            elif idx >= len(loc_times_sorted):
                best = len(loc_times_sorted) - 1
            else:
                # 比较前一帧和后一帧哪个更接近
                if abs(img_ts - loc_times_sorted[idx-1]) <= abs(loc_times_sorted[idx] - img_ts):
                    best = idx - 1
                else:
                    best = idx

            pose_flat = loc_poses_sorted[best]
            # line = "{:.6f}".format(img_ts) + " " + " ".join("{:.6f}".format(v) for v in pose_flat)
            line = " ".join("{:.6f}".format(v) for v in pose_flat)
            fout.write(line + "\n")

    print(f"  >> 已生成对齐后文件: {output_txt}")

def main():
    parser = argparse.ArgumentParser(
        description="将 /ins/localization 里提取的 pose 与 /camera/image_left 话题的时间戳对齐，"
                    "给每帧图像分配一个最近邻 pose。"
    )
    parser.add_argument(
        "--bag", "-b", required=True,
        help="待处理的 rosbag 文件路径，例如：my_data.bag"
    )
    parser.add_argument(
        "--output", "-o", required=True,
        help="输出对齐后的文本文件路径，每行：<img_ts> + 12 个 pose 值"
    )
    parser.add_argument(
        "--euler_deg", action="store_true",
        help="如果 /ins/localization 里 roll/pitch/azimuth 单位是度，则加此标志。"
    )

    args = parser.parse_args()
    bag_path = args.bag

    # 1) 提取所有 localization pose
    loc_times, loc_poses = extract_localization_poses(bag_path, euler_in_degrees=args.euler_deg)

    if len(loc_times) == 0:
        print("[Error] 没有从 /ins/localization 中读取到任何消息！请检查话题名称或 msg 字段。")
        return

    # 2) 读取所有图像时间戳并做对齐
    align_with_images(bag_path, loc_times, loc_poses, args.output)


if __name__ == "__main__":
    main()
